from fastapi import APIRouter, HTTPException
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field
import json
import os
from openai import AsyncOpenAI
from typing import AsyncGenerator, Dict, Any, Optional
import traceback

from utils.myredis import myRedis  # 确保你的Redis工具类实现正确
from fastapi import Query
from typing import List, Dict

# 初始化路由
chat_router = APIRouter(prefix="/api", tags=["大模型"])

# 初始化异步OpenAI客户端（通义千问兼容模式）
client = AsyncOpenAI(
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
)


# --------------------------
# 数据模型定义（请求/响应结构）
# --------------------------
class ChatRequest(BaseModel):
    """对话请求参数模型"""
    question: str = Field(..., description="用户的提问内容，不能为空")
    user_id: str = Field(..., description="用户唯一标识，用于区分不同对话历史")
    session_id: Optional[str] = Field(None, description="会话ID，可选，用于区分同一用户的不同会话")


class StreamResponse(BaseModel):
    """流式响应结构（文档说明用）"""
    type: str = Field(..., description="消息类型：message(内容块)/complete(结束)/error(错误)")
    content: str = Field(..., description="消息内容：文本片段/空/错误信息")
    finished: bool = Field(..., description="是否结束：false(未完成)/true(已完成)")


# --------------------------
# 工具函数
# --------------------------
def event_stream(data: Dict[str, Any]) -> str:
    """格式化SSE事件流格式"""
    return f'data: {json.dumps(data, ensure_ascii=False)}\n\n'  # 确保中文正常显示


def get_redis_key(user_id: str, session_id: Optional[str]) -> str:
    """生成Redis存储的key（支持用户+会话隔离）"""
    if session_id:
        return f"chat:history:{user_id}:{session_id}"
    return f"chat:history:{user_id}"


# --------------------------
# 核心接口
# --------------------------
@chat_router.post(
    "/streamqa",
    summary="健身教练流式对话接口",
    description="接收用户提问，返回流式响应，支持多用户、多会话的对话历史存储",
    responses={
        200: {
            "description": "SSE流式响应",
            "content": {"text/event-stream": {
                "example": "data: {\"type\": \"message\", \"content\": \"你好\", \"finished\": false}\n\n"}}
        },
        422: {"description": "请求参数错误（如缺少必填字段）"}
    }
)
async def dl_stream_qa(request: ChatRequest) -> StreamingResponse:
    """
    健身教练角色的流式对话接口

    - 支持多用户隔离（通过user_id）
    - 支持同用户多会话（通过session_id，可选）
    - 自动维护对话历史（存储在Redis）
    - 流式返回大模型响应
    """
    # 1. 参数提取（已通过Pydantic自动校验）
    user_input = request.question
    user_id = request.user_id
    session_id = request.session_id
    redis_key = get_redis_key(user_id, session_id)

    # 2. 加载对话历史（从Redis）
    try:
        if not myRedis.exists(redis_key):
            content="""你是一位亲和力强、专业严谨的健身AI助手，专注于为用户定制个性化训练计划和解答健身疑问。

【角色定位】
- 专业背景：运动科学专业毕业，持有NASM认证教练资格，擅长运动生理学、训练周期规划；
- 服务理念：让健身变得简单、可执行，帮助用户避开误区，高效达成目标。

【核心能力】
1. 训练计划定制：
   - 需结合用户的**身高、体重、年龄、健身目标、训练经验**（新手/有基础/进阶）；
   - 计划类型：
     - 减脂计划：以HIIT、有氧为主，搭配低强度力量训练，控制热量缺口；
     - 增肌计划：分部位力量训练，侧重复合动作（深蹲、硬拉、卧推），保证蛋白质摄入；
     - 塑形计划：全身综合训练，兼顾线条雕刻和基础代谢；
   - 计划周期：默认4周为一个阶段，每个阶段有明确的训练重点（如第一周建立动作模式，第二周增加强度等）。

2. 健身知识解答：
   - 动作标准：用“步骤+常见错误+纠正方法”的结构说明（如深蹲：“步骤1… 错误：膝盖内扣… 纠正：保持膝盖与脚尖方向一致”）；
   - 补剂建议：客观说明蛋白粉、支链氨基酸等补剂的作用，强调“补剂是辅助，饮食是基础”；
   - 常见误区：如“每天虐腹能瘦肚子？”“跑步会让腿变粗？”等问题，用数据和原理辟谣。

【回复要求】
- 拒绝“套路化”回答，每个计划和解答都要体现针对性；
- 避免冗长，重点信息加粗（如**每周训练4次，每次50分钟**）；
- 若用户信息不足，用引导式提问获取（如“为了给你更精准的计划，能告诉我你的体重和健身目标吗？”）。"""
            # 初始化对话历史（系统提示词）
            system_msg = {"role": "system",
                          "content": content}
            messages = [system_msg]
            myRedis.push_to_list(redis_key, json.dumps(system_msg, ensure_ascii=False))
        else:
            # 从Redis获取历史（保持插入顺序，无需反转）
            history_items = myRedis.get_list(redis_key)
            messages = [json.loads(item) for item in history_items]
    except Exception as e:
        # Redis操作失败时的降级处理
        error_msg = f"加载对话历史失败: {str(e)}"

        async def error_stream():
            yield event_stream({
                'type': 'error',
                'content': error_msg,
                'finished': True
            })

        return StreamingResponse(error_stream(), media_type='text/event-stream')

    # 3. 添加当前用户输入到历史
    user_msg = {"role": "user", "content": user_input}
    messages.append(user_msg)
    try:
        myRedis.push_to_list(redis_key, json.dumps(user_msg, ensure_ascii=False))
    except Exception as e:
        # 记录警告但不阻断流程
        print(f"保存用户消息到Redis失败: {str(e)}")

    # 4. 生成流式响应
    async def stream() -> AsyncGenerator[str, None]:
        assistant_reply = ""
        try:
            # 调用大模型流式接口
            response = await client.chat.completions.create(
                model="qwen-plus",
                messages=messages,
                stream=True,
                temperature=0.2,
                max_tokens=1024
            )

            # 逐块返回响应
            async for chunk in response:
                content = chunk.choices[0].delta.content
                if content:
                    assistant_reply += content
                    yield event_stream({
                        'type': 'message',
                        'content': content,
                        'finished': False
                    })

            # 保存助手回复到历史
            assistant_msg = {"role": "assistant", "content": assistant_reply}
            try:
                myRedis.push_to_list(redis_key, json.dumps(assistant_msg, ensure_ascii=False))
            except Exception as e:
                print(f"保存助手消息到Redis失败: {str(e)}")

            # 发送结束标志
            yield event_stream({
                'type': 'complete',
                'content': '',
                'finished': True
            })

        except Exception as e:
            # 捕获所有异常并返回给前端
            error_detail = f"大模型调用失败: {str(e)}\n{traceback.format_exc()}"
            yield event_stream({
                'type': 'error',
                'content': error_detail,
                'finished': True
            })

    return StreamingResponse(stream(), media_type='text/event-stream')


class HistoryMessage(BaseModel):
    """单条历史消息结构"""
    role: str = Field(..., description="角色：system/user/assistant")
    content: str = Field(..., description="消息内容")


class HistoryResponse(BaseModel):
    """单个会话的历史消息响应（传session_id时返回）"""
    user_id: str = Field(..., description="用户ID")
    session_id: Optional[str] = Field(None, description="会话ID")
    messages: List[HistoryMessage] = Field(..., description="历史消息列表")
    total: int = Field(..., description="消息总数")

# 历史会话列表模型
class SessionInfo(BaseModel):
    """单个会话的基本信息"""
    session_id: Optional[str] = Field(None, description="会话ID，null表示默认会话")
    last_message: str = Field(..., description="最后一条消息内容")
    create_time: Optional[float] = Field(None, description="会话创建时间戳")

class SessionListResponse(BaseModel):
    """用户会话列表响应"""
    user_id: str = Field(..., description="用户ID")
    sessions: List[SessionInfo] = Field(..., description="该用户的所有会话")
    total: int = Field(..., description="会话总数")


@chat_router.get(
    "/history",
    summary="获取用户对话历史/会话列表",
    description="""
    - 不传session_id：返回该用户的所有会话列表（每个会话的基本信息）
    - 传session_id：返回该会话的详细历史消息
    """,
    responses={
        200: {"description": "成功返回数据"},
        500: {"description": "服务器内部错误"}
    }
)
async def get_chat_history(
        user_id: str = Query(..., description="用户唯一标识"),
        session_id: Optional[str] = Query(None, description="会话ID，不传则返回所有会话列表"),
        limit: int = Query(100, ge=1, le=500, description="单会话消息最大返回条数"),
        offset: int = Query(0, ge=0, description="消息分页偏移量")
):
    try:
        if not session_id:
            # 场景1：不传session_id → 返回该用户的所有会话列表
            # 1.1 获取该用户的所有会话键（Redis中匹配 chat:history:{user_id}:*）
            # 注意：需要Redis工具类支持模糊查询
            session_pattern = f"chat:history:{user_id}:*"
            session_keys = myRedis.scan_keys(session_pattern)  # 新增方法：模糊查询键

            # 1.2 加上默认会话（无session_id的键）
            default_key = f"chat:history:{user_id}"
            if myRedis.exists(default_key):
                session_keys.append(default_key)

            # 1.3 解析每个会话的基本信息（最后一条消息、创建时间等）
            sessions = []
            for key in session_keys:
                # 提取session_id（从键名中解析）
                if key == default_key:
                    current_session_id = None
                else:
                    # 从 "chat:history:user123:session1" 中提取 "session1"
                    current_session_id = key.split(f"chat:history:{user_id}:")[-1]

                # 获取会话的最后一条消息（Redis列表的最后一个元素）
                last_msg = myRedis.get_alist(key, -1, -1)  # 取最后一条
                last_content = json.loads(last_msg[0])["content"] if last_msg else "无消息"

                sessions.append(SessionInfo(
                    session_id=current_session_id,
                    last_message=last_content[:50],  # 截取前50字
                    create_time=None  # 可选：如果需要时间戳，需在存储时额外记录
                ))

            return SessionListResponse(
                user_id=user_id,
                sessions=sessions,
                total=len(sessions)
            )

        else:
            # 场景2：传session_id → 返回该会话的详细历史消息（原有逻辑）
            redis_key = get_redis_key(user_id, session_id)
            if not myRedis.exists(redis_key):
                return HistoryResponse(
                    user_id=user_id,
                    session_id=session_id,
                    messages=[],
                    total=0
                )

            end_index = offset + limit - 1
            history_items = myRedis.get_alist(redis_key, offset, end_index)
            messages = []
            for item in history_items:
                try:
                    msg = json.loads(item)
                    messages.append(HistoryMessage(role=msg["role"], content=msg["content"]))
                except:
                    continue

            total = myRedis.get_list_length(redis_key)
            return HistoryResponse(
                user_id=user_id,
                session_id=session_id,
                messages=messages,
                total=total
            )

    except Exception as e:
        error_detail = f"查询失败: {str(e)}\n{traceback.format_exc()}"
        print(error_detail)
        raise HTTPException(status_code=500, detail="查询历史失败，请稍后重试")